//
// This Stan program defines a simple model, with a
// vector of values 'y' modeled as normally distributed
// with mean 'mu' and standard deviation 'sigma'.
//
// Learn more about model development with Stan at:
//
//    http://mc-stan.org/users/interfaces/rstan.html
//    https://github.com/stan-dev/rstan/wiki/RStan-Getting-Started
//

// heterocedastic model
data {
  int<lower=1> N;          // Number of observations
  vector[N] t1;            // Predictor (continuous)
  int<lower=0,upper=1> pro[N];  // Group indicator (0 = non-pro, 1 = pro)
  vector[N] t2;            // Outcome
  vector[N] weights;       // Weights for each observation

}

parameters {
  real beta_0;             // Intercept
  real beta_1;             // Effect of t1
  real beta_2;             // Effect of pro
  real beta_3;             // Interaction t1 * pro
  
  real<lower=0> sigma_nonpro;  // SD for non-professionals
  real<lower=0> sigma_pro;     // SD for professionals
}

model {
  // Priors
  beta_0 ~ normal(0, 1);
  beta_1 ~ normal(0, 1);
  beta_2 ~ normal(0, 1);
  beta_3 ~ normal(0, 1);
  
  sigma_nonpro ~ normal(0, 1);
  sigma_pro ~ normal(0, 1);
  
  // Likelihood
  for (i in 1:N) {
    real sigma = pro[i] * sigma_pro + (1 - pro[i]) * sigma_nonpro; 
    target += weights[i] * normal_lpdf(t2[i] | beta_0 + beta_1 * t1[i] + beta_2 * pro[i] + beta_3 * t1[i] * pro[i], sigma);
  }
}
